package com.ruoyi.nsh.service.impl;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Service;
import com.ruoyi.nsh.service.SqlQueryService;
import com.ruoyi.framework.datasource.DynamicDataSourceContextHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.ruoyi.common.core.page.PageDomain;
import com.ruoyi.common.core.page.TableSupport;

import java.util.List;
import java.util.Map;
import java.util.ArrayList;
import java.util.stream.Collectors;

@Service
public class SqlQueryServiceImpl implements SqlQueryService {
    
    private static final Logger log = LoggerFactory.getLogger(SqlQueryServiceImpl.class);
    
    @Autowired
    private JdbcTemplate jdbcTemplate;
    
    @Override
    public List<Map<String, Object>> executeQuery(String sql, String dataSource) {
        try {
            DynamicDataSourceContextHolder.setDataSourceType(dataSource);
            log.info("Using {} DataSource", dataSource);
            
            PageDomain pageDomain = TableSupport.buildPageRequest();
            Integer pageNum = pageDomain.getPageNum();
            Integer pageSize = pageDomain.getPageSize();
            
            String paginatedSql;
            if ("SLAVE".equals(dataSource)) {  // DB2
                // 先获取原始查询的列名和顺序
                String columnSql = String.format("SELECT * FROM (%s) AS TEMP FETCH FIRST 1 ROWS ONLY WITH UR", sql);
                List<Map<String, Object>> firstRow = jdbcTemplate.queryForList(columnSql);
                
                if (!firstRow.isEmpty()) {
                    // 保持原始列的顺序
                    List<String> columnList = new ArrayList<>(firstRow.get(0).keySet());
                    String columns = columnList.stream()
                        .map(col -> "TEMP." + col)  // 添加表别名前缀
                        .collect(Collectors.joining(", "));
                    
                    paginatedSql = String.format(
                        "WITH NUMBERED_ROWS AS (" +
                        "  SELECT %s, ROW_NUMBER() OVER() AS RN FROM (%s) AS TEMP" +
                        ") " +
                        "SELECT %s FROM NUMBERED_ROWS " +
                        "WHERE RN BETWEEN %d AND %d",
                        columns,
                        sql,
                        columns.replace("TEMP.", ""),  // 移除最终结果中的表别名
                        (pageNum - 1) * pageSize + 1,
                        pageNum * pageSize
                    );
                } else {
                    paginatedSql = sql;
                }
            } else {  // MySQL或GBase
                paginatedSql = String.format("%s LIMIT %d OFFSET %d",
                    sql,
                    pageSize,
                    (pageNum - 1) * pageSize
                );
            }
            
            List<Map<String, Object>> result = jdbcTemplate.queryForList(paginatedSql);
            
            // 对于DB2，移除分页相关的字段
            if ("SLAVE".equals(dataSource)) {
                result.forEach(row -> row.remove("RN"));
            }
            
            return result;
        } finally {
            DynamicDataSourceContextHolder.clearDataSourceType();
        }
    }

    @Override
    public long queryTotal(String sql, String dataSource) {
        try {
            DynamicDataSourceContextHolder.setDataSourceType(dataSource);
            String countSql = "SELECT COUNT(1) as total FROM (" + sql + ") t";
            Long total = jdbcTemplate.queryForObject(countSql, Long.class);
            return total != null ? total : 0L;
        } finally {
            DynamicDataSourceContextHolder.clearDataSourceType();
        }
    }

    @Override
    public List<Map<String, Object>> getTableList(String dataSource) {
        try {
            DynamicDataSourceContextHolder.setDataSourceType(dataSource);
            String sql;
            if ("SLAVE".equals(dataSource)) {  // DB2
                sql = "SELECT TRIM(t.TABSCHEMA) as SCHEMA_NAME, " +
                      "TRIM(t.TABNAME) as TABLENAME, " +
                      "TRIM(t.REMARKS) as TABLE_COMMENT, " +
                      "CARD as ROW_COUNT, " +
                      "CAST(FLOAT(NPAGES) * 4096 / 1024 / 1024 as DECIMAL(10,2)) as SIZE_MB " +
                      "FROM SYSCAT.TABLES t " +
                      "WHERE t.TABSCHEMA NOT IN ('SYSCAT', 'SYSIBM', 'SYSSTAT', 'SYSTOOLS')";
            } else if ("GBASE".equals(dataSource)) {
                sql = "SELECT t.TABLE_SCHEMA as SCHEMA_NAME, " +
                      "t.TABLE_NAME as TABLENAME, " +
                      "t.TABLE_COMMENT as TABLE_COMMENT, " +
                      "t.TABLE_ROWS as ROW_COUNT, " +
                      "ROUND((t.DATA_LENGTH + t.INDEX_LENGTH) / 1024 / 1024, 2) as SIZE_MB " +
                      "FROM INFORMATION_SCHEMA.TABLES t " +
                      "WHERE t.TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')";
            } else {  // MySQL
                sql = "SELECT t.TABLE_SCHEMA as SCHEMA_NAME, " +
                      "t.TABLE_NAME as TABLENAME, " +
                      "t.TABLE_COMMENT as TABLE_COMMENT, " +
                      "t.TABLE_ROWS as ROW_COUNT, " +
                      "ROUND((t.DATA_LENGTH + t.INDEX_LENGTH) / 1024 / 1024, 2) as SIZE_MB " +
                      "FROM INFORMATION_SCHEMA.TABLES t " +
                      "WHERE t.TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')";
            }
            return jdbcTemplate.queryForList(sql);
        } finally {
            DynamicDataSourceContextHolder.clearDataSourceType();
        }
    }

    @Override
    public List<String> getSchemaList(String dataSource) {
        try {
            DynamicDataSourceContextHolder.setDataSourceType(dataSource);
            String sql;
            if ("SLAVE".equals(dataSource)) {  // DB2
                sql = "SELECT DISTINCT TRIM(SCHEMANAME) FROM SYSCAT.SCHEMATA " +
                      "WHERE DEFINER != 'SYSIBM' AND SCHEMANAME NOT IN ('SYSCAT', 'SYSIBM', 'SYSSTAT', 'SYSTOOLS')";
            } else if ("GBASE".equals(dataSource)) {
                sql = "SELECT DISTINCT TABLE_SCHEMA FROM INFORMATION_SCHEMA.TABLES " +
                      "WHERE TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')";
            } else {  // MySQL
                sql = "SELECT DISTINCT TABLE_SCHEMA FROM INFORMATION_SCHEMA.TABLES " +
                      "WHERE TABLE_SCHEMA NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys')";
            }
            return jdbcTemplate.queryForList(sql, String.class);
        } finally {
            DynamicDataSourceContextHolder.clearDataSourceType();
        }
    }
} 